#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
@File        : mngs_replace_ref.py
@Author      : Bing Liang
@Email       : believer19940901@gmail.com
@Date        : 2025/11/04 14:33
@Description :
    替换或新增 pan-genome 数据库中指定物种的参考序列。

    功能：
        - 如果 pan-genome 中已存在该 species_id，则替换其序列；
        - 如果不存在，则新增该序列。

    适用场景：
        用于更新某物种的代表参考序列，或者向 pan-genome 中添加新物种。
"""

from argparse import ArgumentParser, Namespace
from pathlib import Path
from Bio import SeqIO


def mngs_replace_ref(args: Namespace):
    """
    替换或新增物种参考序列。
    """

    pan_genome: Path = args.pan_genome.resolve()
    ref_fasta: Path = args.ref_fasta.resolve()
    out_file: Path = args.out_file.resolve()
    species_id: str = args.species_id.strip()

    # 确保输出目录存在
    out_file.parent.mkdir(parents=True, exist_ok=True)

    # 读取新的参考序列（只取第一条）
    try:
        ref_record = next(SeqIO.parse(ref_fasta, "fasta"))
    except StopIteration:
        raise ValueError(f"参考序列文件 {ref_fasta} 中未找到任何序列。")

    # 标准化序列 ID/描述（保持 pan-genome 中的 ID 风格）
    ref_record.id = species_id
    ref_record.name = species_id
    ref_record.description = species_id

    out_records = []
    replaced = False

    # 遍历原 pan-genome，替换或保留
    for record in SeqIO.parse(pan_genome, "fasta"):
        if record.id == species_id:
            # 找到目标物种 → 替换
            out_records.append(ref_record)
            replaced = True
        else:
            out_records.append(record)

    # 如果未替换，说明 pan-genome 中没有该物种 → 追加
    if not replaced:
        out_records.append(ref_record)

    # 写出结果
    SeqIO.write(out_records, out_file, "fasta")

    # 输出提示信息
    if replaced:
        print(f"[OK] 已替换物种 '{species_id}' 的参考序列 → {out_file}")
    else:
        print(f"[OK] 未在 pan-genome 中找到 '{species_id}'，已执行新增 → {out_file}")


def main():
    parser = ArgumentParser(description="替换或新增 pan-genome 中的物种参考序列")
    parser.add_argument("--pan_genome", type=Path, required=True, help="原始 pan-genome FASTA 文件路径")
    parser.add_argument("--ref_fasta", type=Path, required=True, help="新的参考序列 FASTA 文件路径（只取第一条）")
    parser.add_argument("--out_file", type=Path, required=True, help="输出文件路径")
    parser.add_argument("--species_id", type=str, required=True, help="物种 ID，例如 '1280_Staphylococcus_aureus'")
    args = parser.parse_args()
    mngs_replace_ref(args)


if __name__ == "__main__":
    main()
